"""
predict.py
单张图片推理脚本
使用：python predict.py --img 枸杞001.jpg --weights models/best.pth
"""
import argparse, torch
from PIL import Image
from src.model import get_model
from src.dataset import HerbDataset
from src.config import Config

def predict(img_path, weights):
    # 加载模型
    model = get_model().to(Config.device)
    model.load_state_dict(torch.load(weights, map_location=Config.device))
    model.eval()

    # 复用 Dataset 里的预处理
    trans = HerbDataset(Config.test_dir, Config.input_size, 'test').trans

    # 读取并预处理图片
    img = Image.open(img_path).convert('RGB')
    img_tensor = trans(img).unsqueeze(0).to(Config.device)  # 增加 batch 维度

    # 前向推理
    with torch.no_grad():
        outputs = torch.softmax(model(img_tensor), dim=1)  # 转概率
        prob, idx = torch.max(outputs, 1)  # 取最大概率和对应类别索引

    # 类别名称
    classes = ['猫', '狗', '其他']
    print(f'✅预测结果：{classes[idx]}  置信度：{prob.item():.2%}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img', required=True, help='待预测图片路径')
    parser.add_argument('--weights', default='models/best.pth', help='模型权重路径')
    args = parser.parse_args()
    predict(args.img, args.weights)